#define NOMINMAX

#include <algorithm/analytical_model/algorithm/cuda/Utils.h>
#include <algorithm/analytical_model/algorithm/cuda/VectorMath.cuh>
#include <algorithm/analytical_model/algorithm/cuda/grid/RectangleGrid.cuh>
#include <algorithm/analytical_model/algorithm/cuda/heliostat/RectangleHeliostat.cuh>
#include <utils/BitUtils.h>
#include <utils/cuda/algorithms/ParallelScan.cuh>
#include <utils/cuda/containers/Vector.cuh>

#include <cuda_device_runtime_api.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <driver_types.h>

#include <thrust/async/for_each.h>
#include <thrust/async/scan.h>
#include <thrust/async/transform.h>
#include <thrust/device_ptr.h>
#include <thrust/functional.h>
#include <thrust/iterator/counting_iterator.h>

#include <algorithm>
#include <memory>
#include <tuple>

using namespace solar;
using namespace solar::cuda;

void RectangleGrid::init()
{
    subgrid_num_.x = int(std::ceil(double(size_.x / interval_.x)));
    subgrid_num_.y = int(std::ceil(double(size_.y / interval_.y)));
    subgrid_num_.z = int(std::ceil(double(size_.z / interval_.z)));
    CHECK_CUDA_ERRORS(cudaStreamCreate(&stream_));
    CHECK_CUDA_ERRORS(cudaHostAlloc(std::addressof(h_num_grid_helio_match_ptr_), sizeof(int) * 1,
                                    cudaHostAllocMapped));
    CHECK_CUDA_ERRORS(cudaHostAlloc(std::addressof(h_d_grid_helio_match_ptr_), sizeof(int*) * 1,
                                    cudaHostAllocMapped));
    CHECK_CUDA_ERRORS(cudaHostGetDevicePointer(std::addressof(d_num_grid_helio_match_ptr_),
                                               h_num_grid_helio_match_ptr_, 0));
    CHECK_CUDA_ERRORS(cudaHostGetDevicePointer(std::addressof(d_d_grid_helio_match_ptr_),
                                               h_d_grid_helio_match_ptr_, 0));

    // printf("%d, %d, %d\n", subgrid_num_.x, subgrid_num_.y, subgrid_num_.z);
}

void RectangleGrid::clear()
{
    if (d_grid_helio_index_ != nullptr)
    {
        CHECK_CUDA_ERRORS(cudaFreeAsync(d_grid_helio_index_, stream_));
        d_grid_helio_index_ = nullptr;
    }

    if (d_grid_helio_match_ != nullptr)
    {
        // CHECK_CUDA_ERRORS(cudaFreeAsync(d_grid_helio_match_, stream_));
        auto* d_d_grid_helio_match_ptr = this->d_d_grid_helio_match_ptr_;
        thrust::for_each(thrust::device.on(stream_), thrust::make_counting_iterator(0),
                         thrust::make_counting_iterator(1),
                         [d_d_grid_helio_match_ptr] __device__(const int&)
                         { free(*d_d_grid_helio_match_ptr); });
        d_grid_helio_match_ = nullptr;
    }
    if (h_num_grid_helio_match_ptr_ != nullptr)
    {
        CHECK_CUDA_ERRORS(cudaFreeHost(h_num_grid_helio_match_ptr_));
        h_num_grid_helio_match_ptr_ = nullptr;
    }
    if (h_d_grid_helio_match_ptr_ != nullptr)
    {
        CHECK_CUDA_ERRORS(cudaFreeHost(h_d_grid_helio_match_ptr_));
        h_d_grid_helio_match_ptr_ = nullptr;
    }
    cudaStreamDestroy(stream_);
}

auto RectangleGrid::gridHelioMatch(RectangleHeliostat* d_heliostats) -> void
{
    // thrust::exclusive_scan 和 thrust::inclusive_scan似乎不能正常使用？

    // if (d_grid_helio_match_ || d_grid_helio_index_)
    // {
    //     throw std::runtime_error("The grid and heliostats corresponding relationship should be "
    //                              "empty before calling this method.");
    // }

    thrust::device_ptr<RectangleHeliostat> heliostats_ptr(d_heliostats + start_helio_index_);

    int* start_subhelio_pos_vec_raw_raw_ptr = nullptr;
    int* start_subhelio_pos_vec_raw_ptr = nullptr;
    Vector<std::tuple<int, int, int, int>>* grid_helio_match_vector_per_heliostat_raw_ptr = nullptr;
    CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(start_subhelio_pos_vec_raw_raw_ptr),
                                      sizeof(int) * num_helios_, stream_));
    CHECK_CUDA_ERRORS(
        cudaMemsetAsync(start_subhelio_pos_vec_raw_raw_ptr, 0, sizeof(int) * num_helios_, stream_));
    CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(start_subhelio_pos_vec_raw_ptr),
                                      sizeof(int) * num_helios_, stream_));
    CHECK_CUDA_ERRORS(
        cudaMemsetAsync(start_subhelio_pos_vec_raw_ptr, 0, sizeof(int) * num_helios_, stream_));
    CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(grid_helio_match_vector_per_heliostat_raw_ptr),
                                      sizeof(Vector<std::tuple<int, int, int, int>>) * num_helios_,
                                      stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(grid_helio_match_vector_per_heliostat_raw_ptr, 0,
                                      sizeof(int) * num_helios_, stream_));
    thrust::device_ptr<int> start_subhelio_pos_vec_raw(start_subhelio_pos_vec_raw_raw_ptr);
    thrust::device_ptr<int> start_subhelio_pos_vec(start_subhelio_pos_vec_raw_ptr);
    thrust::device_ptr<Vector<std::tuple<int, int, int, int>>>
        grid_helio_match_vector_per_heliostat(grid_helio_match_vector_per_heliostat_raw_ptr);
    // cudaStreamSynchronize(stream_);
    thrust::async::transform(thrust::device.on(stream_), heliostats_ptr,
                             heliostats_ptr + num_helios_, start_subhelio_pos_vec_raw,
                             [] __device__(const RectangleHeliostat& h1)
                             { return h1.sub_helio_size_; });
    // cudaStreamSynchronize(stream_);

    // Actually it should be exclusiveScan with init = 0
    // algorithm::inclusiveScan(start_subhelio_pos_vec_raw_raw_ptr + 1,
    //                          start_subhelio_pos_vec_raw_ptr + 1, num_helios_ - 1, stream_);
    thrust::async::inclusive_scan(thrust::device.on(stream_),
                                  start_subhelio_pos_vec_raw_raw_ptr + 1,
                                  start_subhelio_pos_vec_raw_raw_ptr + num_helios_,
                                  start_subhelio_pos_vec_raw_ptr + 1, thrust::plus<>());
    CHECK_CUDA_ERRORS(cudaGetLastError());
    // cudaStreamSynchronize(stream_);

    const float3 pos = this->pos_;
    const float3 interval = this->interval_;
    const int3 subgrid_num = this->subgrid_num_;

    auto start_helio_index = start_helio_index_;
    // std::cout << start_helio_index << '\n';

    thrust::async::for_each(
        thrust::device.on(stream_), thrust::make_counting_iterator(0),
        thrust::make_counting_iterator(num_helios_),
        [pos, interval, subgrid_num, start_helio_index, d_heliostats,
         start_subhelio_pos_vec_raw_ptr,
         grid_helio_match_vector_per_heliostat_raw_ptr] __device__(const int& index)
        {
            auto& heliostat = d_heliostats[start_helio_index + index];
            auto& current_start_subhelio_pos = start_subhelio_pos_vec_raw_ptr[index];
            grid_helio_match_vector_per_heliostat_raw_ptr[index] =
                Vector<std::tuple<int, int, int, int>>();
            const auto heliostat_pos = heliostat.pos_;
            const auto radius = length(heliostat.size_) / 2;
            const auto min_pos = heliostat_pos - radius;
            const auto max_pos = heliostat_pos + radius;

            // boxIntersect

            const int3 min_grid_pos = make_int3(floor((min_pos - pos) / interval));
            const int3 max_grid_pos = make_int3(floor((max_pos - pos) / interval));

            for (int x = std::max(0, min_grid_pos.x);
                 x <= std::min(subgrid_num.x - 1, max_grid_pos.x); ++x)
            {
                for (int y = std::max(0, min_grid_pos.y);
                     y <= std::min(subgrid_num.y - 1, max_grid_pos.y); ++y)
                {
                    for (int z = std::max(0, min_grid_pos.z);
                         z <= std::min(subgrid_num.z - 1, max_grid_pos.z); ++z)
                    {
                    // int pos_id = (x * subgrid_num.y * subgrid_num.z) + (y
                    // * subgrid_num.z) + z;

#pragma unroll
                        for (int i = 0; i < heliostat.sub_helio_size_; ++i)
                        {
                            grid_helio_match_vector_per_heliostat_raw_ptr[index].push_back(
                                {x, y, z, current_start_subhelio_pos + i});
                        }
                    }
                }
            }
        });

    // cudaStreamSynchronize(stream_);
    // std::cout << "for_each match\n";

    CHECK_CUDA_ERRORS(cudaMallocAsync(
        std::addressof(d_grid_helio_index_),
        sizeof(int) * subgrid_num_.x * subgrid_num_.y * subgrid_num_.z + 2, stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(d_grid_helio_index_, 0, sizeof(int) * 1, stream_));
    // cudaStreamSynchronize(stream_);
    thrust::device_ptr<int> d_grid_helio_index_ptr(d_grid_helio_index_);

    Vector<int>* grid_helio_match_vector_per_grid_raw_ptr = nullptr;
    CHECK_CUDA_ERRORS(cudaMallocAsync(
        std::addressof(grid_helio_match_vector_per_grid_raw_ptr),
        sizeof(Vector<int>) * subgrid_num_.x * subgrid_num_.y * subgrid_num_.z, stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(
        grid_helio_match_vector_per_grid_raw_ptr, 0,
        sizeof(Vector<int>) * subgrid_num_.x * subgrid_num_.y * subgrid_num_.z, stream_));
    thrust::device_ptr<Vector<int>> grid_helio_match_vector_per_grid_ptr(
        grid_helio_match_vector_per_grid_raw_ptr);
    thrust::async::for_each(thrust::device.on(stream_), grid_helio_match_vector_per_grid_ptr,
                            grid_helio_match_vector_per_grid_ptr +
                                subgrid_num_.x * subgrid_num_.y * subgrid_num_.z,
                            []__device__(Vector<int>& vec) { vec = Vector<int>(); });
    // std::cout << "for_each grid_helio_match_vector_per_grid init\n";

    int* grid_helio_match_vector_per_grid_size_raw_ptr = nullptr;
    // int* grid_helio_match_vector_per_grid_size_raw_ptr2 = nullptr;
    int* grid_helio_match_vector_per_grid_offset_raw_ptr = d_grid_helio_index_;
    int* grid_helio_match_vector_per_grid_current_index_raw_ptr = nullptr;
    CHECK_CUDA_ERRORS(cudaMallocAsync(
        std::addressof(grid_helio_match_vector_per_grid_size_raw_ptr),
        sizeof(int) * this->subgrid_num_.x * this->subgrid_num_.y * this->subgrid_num_.z, stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(
        grid_helio_match_vector_per_grid_size_raw_ptr, 0,
        sizeof(int) * this->subgrid_num_.x * this->subgrid_num_.y * this->subgrid_num_.z, stream_));
    // CHECK_CUDA_ERRORS(cudaMallocAsync(
    //     std::addressof(grid_helio_match_vector_per_grid_size_raw_ptr2),
    //     sizeof(int) * this->subgrid_num_.x * this->subgrid_num_.y * this->subgrid_num_.z,
    //     stream_));
    CHECK_CUDA_ERRORS(cudaMallocAsync(
        std::addressof(grid_helio_match_vector_per_grid_current_index_raw_ptr),
        sizeof(int) * this->subgrid_num_.x * this->subgrid_num_.y * this->subgrid_num_.z, stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(
        grid_helio_match_vector_per_grid_current_index_raw_ptr, 0,
        sizeof(int) * this->subgrid_num_.x * this->subgrid_num_.y * this->subgrid_num_.z, stream_));
    // cudaStreamSynchronize(stream_);
    thrust::device_ptr<int> grid_helio_match_vector_per_grid_size_ptr(
        grid_helio_match_vector_per_grid_size_raw_ptr);
    thrust::async::for_each(
        thrust::device.on(stream_), grid_helio_match_vector_per_heliostat,
        grid_helio_match_vector_per_heliostat + num_helios_,
        [grid_helio_match_vector_per_grid_size_raw_ptr,
         subgrid_num] __device__(const Vector<std::tuple<int, int, int, int>>& current_match)
        {
#pragma unroll 4
            for (int i = 0; i < current_match.size(); i++)
            {
                const auto& [x, y, z, current_match_result] = current_match[i];
                const int pos_id = (x * subgrid_num.y * subgrid_num.z) + (y * subgrid_num.z) + z;
                atomicAdd(std::addressof(grid_helio_match_vector_per_grid_size_raw_ptr[pos_id]), 1);
            }
        });
    // std::cout << "for_each grid_helio_match_vector_per_grid_size\n";

    // algorithm::inclusiveScan(grid_helio_match_vector_per_grid_size_raw_ptr,
    //                          grid_helio_match_vector_per_grid_offset_raw_ptr + 1,
    //                          this->subgrid_num_.x * this->subgrid_num_.y * this->subgrid_num_.z,
    //                          stream_);
    thrust::async::inclusive_scan(
        thrust::device.on(stream_), grid_helio_match_vector_per_grid_size_raw_ptr,
        grid_helio_match_vector_per_grid_size_raw_ptr +
            this->subgrid_num_.x * this->subgrid_num_.y * this->subgrid_num_.z,
        grid_helio_match_vector_per_grid_offset_raw_ptr, thrust::plus<>());

    CHECK_CUDA_ERRORS(cudaGetLastError());
    // cudaStreamSynchronize(stream_);

    int* d_num_grid_helio_match_ptr = this->d_num_grid_helio_match_ptr_;
    int** d_d_grid_helio_match_ptr = this->d_d_grid_helio_match_ptr_;

    thrust::async::for_each(
        thrust::device.on(stream_), thrust::make_counting_iterator(0),
        thrust::make_counting_iterator(1),
        [d_num_grid_helio_match_ptr, d_d_grid_helio_match_ptr,
         grid_helio_match_vector_per_grid_offset_raw_ptr, subgrid_num] __device__(const int&)
        {
            *d_num_grid_helio_match_ptr =
                grid_helio_match_vector_per_grid_offset_raw_ptr[subgrid_num.x * subgrid_num.y *
                                                                    subgrid_num.z -
                                                                1];
            // printf("num_grid_helio_match=%d\n", *d_num_grid_helio_match_ptr);
            *d_d_grid_helio_match_ptr = new int[*d_num_grid_helio_match_ptr];
            memset(*d_d_grid_helio_match_ptr, 0, sizeof(int) * (*d_num_grid_helio_match_ptr));
        });

    // int num_grid_helio_match = 0;
    // CHECK_CUDA_ERRORS(
    //     cudaMemcpy(std::addressof(num_grid_helio_match),
    //                d_grid_helio_index_ + subgrid_num_.x * subgrid_num_.y * subgrid_num_.z - 1,
    //                sizeof(int) * 1, cudaMemcpyDeviceToHost));
    // num_grid_helio_match_ = num_grid_helio_match;
    // cudaMallocAsync(std::addressof(d_grid_helio_match_), sizeof(int) * num_grid_helio_match_,
    //                 stream_);
    // // cudaStreamSynchronize(stream_);
    // std::cout << num_grid_helio_match_ << '\n';

    // auto* d_grid_helio_match_ptr = d_grid_helio_match_;
    thrust::async::for_each(
        thrust::device.on(stream_), grid_helio_match_vector_per_heliostat,
        grid_helio_match_vector_per_heliostat + num_helios_,
        [grid_helio_match_vector_per_grid_current_index_raw_ptr,
         grid_helio_match_vector_per_grid_offset_raw_ptr, subgrid_num,
         d_d_grid_helio_match_ptr] __device__(const Vector<std::tuple<int, int, int, int>>&
                                                  current_match)
        {
    // printf("%d\n", current_match.size());

#pragma unroll 4
            for (int i = 0; i < current_match.size(); i++)
            {
                const auto& [x, y, z, current_match_result] = current_match[i];
                const int pos_id = (x * subgrid_num.y * subgrid_num.z) + (y * subgrid_num.z) + z;
                const auto current_index = atomicAdd(
                    std::addressof(grid_helio_match_vector_per_grid_current_index_raw_ptr[pos_id]),
                    1);
                const auto current_offset = bit::conditionalAdd2(
                    pos_id >= 1, 0,
                    grid_helio_match_vector_per_grid_offset_raw_ptr[__nv_max(pos_id - 1, 0)]);
                (*d_d_grid_helio_match_ptr)[current_offset + current_index] = current_match_result;
            }
        });
    // std::cout << "for_each data_transform\n";
    CHECK_CUDA_ERRORS(cudaFreeAsync(start_subhelio_pos_vec_raw_raw_ptr, stream_));
    CHECK_CUDA_ERRORS(cudaFreeAsync(start_subhelio_pos_vec_raw_ptr, stream_));
    CHECK_CUDA_ERRORS(cudaFreeAsync(grid_helio_match_vector_per_heliostat_raw_ptr, stream_));
    CHECK_CUDA_ERRORS(cudaFreeAsync(grid_helio_match_vector_per_grid_raw_ptr, stream_));
    CHECK_CUDA_ERRORS(cudaFreeAsync(grid_helio_match_vector_per_grid_size_raw_ptr, stream_));
    CHECK_CUDA_ERRORS(
        cudaFreeAsync(grid_helio_match_vector_per_grid_current_index_raw_ptr, stream_));
}

auto RectangleGrid::gridHelioMatch(const std::vector<RectangleHeliostat*>& heliostats) -> int
{
    return 0;
}

void RectangleGrid::syncCudaStream()
{
    CHECK_CUDA_ERRORS(cudaStreamSynchronize(stream_));
    num_grid_helio_match_ = *h_num_grid_helio_match_ptr_;
    d_grid_helio_match_ = *h_d_grid_helio_match_ptr_;
}
